{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2f2c09d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import display, HTML\n",
    "display(HTML(\"<style>.container { width:80% !important; }</style>\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1352e90",
   "metadata": {},
   "source": [
    "# Intro \n",
    "\n",
    "```\n",
    "This notebook is a minimalistic proof-of-concept showing how a Hailo device can offload the heavy 2D-convolutional part of a point-cloud-input 3D-object-detection network. For this example we use the PointPillars (PP) network from OpenPCDet repo, a modular ecosystem with many 3D networks. The pre/post computation is running in PyTorch using native OpenPCDet code. Almost all of it on CPU, sans the 3D-NMS op which unfortunately is compliled for cuda only.\n",
    "\n",
    "Some wrapping code was necessary to properly integrate openpcdet ot Hailo, see openpcdet2hailo_uitls.py. Some integration details, i.e. how we \"hook\" Hailo into the pytorch model, are net-specific. These semi-manually extracted details are product of an exploration process of the network as given by OpenPCDet.  See the APPENDIX at the bottom of this notebook for an outline of how this has been done, to get a roadmap of how to replicate exploration and adaptation to other (non-PP) 3D network out of OpenPCdet collection, so that to quickly offload ANY openpcdet 3d-detection net to Hailo device.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0bbf447",
   "metadata": {},
   "source": [
    "## Final result\n",
    "Processing a point cloud to get 3D boxes - similar result when executing part of network on Hailo. \n",
    "\n",
    "We'll go through all steps towards this bottom line in the course of this notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e084be19",
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython\n",
    "IPython.display.Image('./rsc/orig_model_front.png', width=480)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "040e0a8d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "IPython.display.Image('./rsc/hailo_assisted_model_front.png', width=480) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5511c37",
   "metadata": {},
   "source": [
    "## OpenPCdet - overview\n",
    "\n",
    "\n",
    "https://github.com/open-mmlab/OpenPCDet\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d73fc49",
   "metadata": {},
   "outputs": [],
   "source": [
    "IPython.display.Image('./rsc/openpcdet_1.png', width=600)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96cb02f9-26e0-4c58-8b47-3fd9899635ba",
   "metadata": {},
   "source": [
    "#### We map to Hailo the 2D backbone and detection head:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4674ebec-0b52-4c6e-988e-ebdfa586805f",
   "metadata": {},
   "outputs": [],
   "source": [
    "IPython.display.Image('./rsc/openpcdet_3.png', width=450)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c79354ac",
   "metadata": {},
   "source": [
    "```\n",
    "The purple bracket marks the part of net we can offload to any Hailo device (\"Hx\" = H8,H15,..). \n",
    "In case of PointPillars, that makes the lion's share of TOPS, as 3D part while complex is not a massive consumer of compute.\n",
    "The 3D parts of this and other nets can be Hailo-mapped too, but with a significant task-specific effort required to achieve good efficiency. Same goes for some postprocessing at the end of \"RPN head\" (anchors2boxes, 3D-NMS) that is also on host yet.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3c7a44e",
   "metadata": {},
   "source": [
    "# Setup\n",
    "\n",
    "## Install, imports, names\n",
    "\n",
    "Run the notebook from a virtualenv prepared like so:\n",
    "\n",
    "1. Install CUDA and Pytorch. Tested configs:\n",
    "    1. `pip install torch==1.12.1+cu113` (assuming CUDA 11.3) **OR**\n",
    "    1. `pip install torch=1.12.1+cu102` (assuming CUDA 10.2)\n",
    "1. Clone & install OpenPCDet: (tested w. commit a68aaa656 04-Apr-23) \n",
    "    ```\n",
    "     pip install -r requirements\n",
    "     pip install spconv kornia\n",
    "     python setup.py develop\n",
    "   ```\n",
    "1. Install Mayavi for 3D visualization of point clouds and 3D boxes. Installing and using Mayavi and its dependencies (PyQT5) might be tricky, especially working remotely on a server, so in the notebook code we create visual results as png files without any windows. Still, it might help to prepend the *jupyter-notebook* launch command with instructions to skip checking for gui toolkit (s.a. Qt) and, if no screen (\"headless\"), also creating and working against a virtual display, like so\n",
    "```\n",
    "ETS_TOOLKIT=null xvfb-run --server-args=\"-screen 0 1024x768x24\" jupyter notebook ...\n",
    "```\n",
    "\n",
    "1. Download pretrained PointPillar pytorch model from the link below into your *openpcdet-clone-location*:\n",
    "[https://drive.google.com/file/d/1wMxWTpU1qUoY3DsCH31WJmvJxcjFXKlm/view?usp=sharing]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ae8a85c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# switch to False to get notebook to run smooth in case of visualization setup issues\n",
    "do_visualize = True   \n",
    "\n",
    "# openpcdet clone path - replace with yours\n",
    "openpcdet_clonedir = '/local/users/alexf/OpenPCDet'  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2bafc41",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pcdet, torch, sys\n",
    "sys.path.append(openpcdet_clonedir+'/tools/')\n",
    "pcdet.__version__, torch.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b74c9a14",
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import glob\n",
    "from pathlib import Path\n",
    "import inspect\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import spconv\n",
    "\n",
    "from pcdet.config import cfg, cfg_from_yaml_file\n",
    "from pcdet.datasets import DatasetTemplate\n",
    "from pcdet.models import build_network, load_data_to_gpu\n",
    "from pcdet.utils import common_utils\n",
    "\n",
    "import os\n",
    "from importlib import reload\n",
    "import openpcdet2hailo_utils as ohu; reload(ohu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ffc1f42",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if do_visualize: \n",
    "    import mayavi.mlab as mlab\n",
    "    from visual_utils import visualize_utils as V    \n",
    "    mlab.init_notebook(backend='png')\n",
    "    mlab.figure()  # test to fail early on mayavi setup issue if any.."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "433a1a67",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "yaml_name = openpcdet_clonedir+'/tools/cfgs/kitti_models/pointpillar.yaml'\n",
    "har_name = 'pp_bev_w_head.har'\n",
    "hef_name = har_name.replace('.har', '.hef')\n",
    "logger = common_utils.create_logger()\n",
    "\n",
    "# (!) pre-trained model:\n",
    "pth_name = openpcdet_clonedir+'/pointpillar_7728.pth'  \n",
    "# ..if missing yet, download manually from:\n",
    "# [https://drive.google.com/file/d/1wMxWTpU1qUoY3DsCH31WJmvJxcjFXKlm/view?usp=sharing]\n",
    "\n",
    "# (!) get sample pointclouds (..running once..) \n",
    "if not os.path.exists('./100'):\n",
    "    !wget https://github.com/supervisely-ecosystem/import-kitti-3d-sample-files/releases/download/v0.0.2/100_test.zip\n",
    "    !unzip 100_test.zip\n",
    "\n",
    "sample_pointclouds = './100/testing/velodyne/'\n",
    "demo_pointcloud = sample_pointclouds+'007000.bin'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7a60812",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_model(cfg, pth_name, demo_dataset):    \n",
    "    model = build_network(model_cfg=cfg.MODEL, num_class=len(cfg.CLASS_NAMES), dataset=demo_dataset)\n",
    "    model.load_params_from_file(filename=pth_name, logger=logger, to_cpu=True)\n",
    "    model.eval()\n",
    "    return model\n",
    "\n",
    "def cfg_from_yaml_file_wrap(yaml_name, cfg):\n",
    "    cwd = os.getcwd()\n",
    "    os.chdir(openpcdet_clonedir+'/tools/')\n",
    "    cfg_from_yaml_file(yaml_name, cfg)\n",
    "    os.chdir(cwd)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07907c83",
   "metadata": {},
   "source": [
    "## Sanity-test: reproduce original OpenPCdet minimalistic demo \n",
    "\n",
    "```  \n",
    "Code based on tools/demo.py as provided by original repo, tweaked a bit to reduce GPU usage to bare minimum.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43d2eef7",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "logger.info('------ Quick Demo of OpenPCDet''s PointPillars -------')\n",
    "cfg_from_yaml_file_wrap(yaml_name, cfg)\n",
    "if False: # \n",
    "    display(cfg.DATA_CONFIG)\n",
    "\n",
    "demo_dataset = ohu.DemoDataset(\n",
    "    dataset_cfg=cfg.DATA_CONFIG, class_names=cfg.CLASS_NAMES, training=False,\n",
    "    root_path=Path(demo_pointcloud), ext='.bin', logger=logger\n",
    ")\n",
    "logger.info(f'Total number of samples: \\t{len(demo_dataset)}')\n",
    "\n",
    "model = get_model(cfg, pth_name, demo_dataset)\n",
    "model_cpu = ohu.PointPillarsCPU(model)\n",
    "\n",
    "with torch.no_grad():\n",
    "    for idx, data_dict in enumerate(demo_dataset):        \n",
    "        data_dict = demo_dataset.collate_batch([data_dict])\n",
    "        pred_dicts, _ = model_cpu(data_dict)        \n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80a4ca0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "if do_visualize: \n",
    "    fig = mlab.figure(size=(1920, 1080))\n",
    "    #fig = mlab.figure(size=(640, 480))\n",
    "    s = ohu.my_draw_scenes(fig=fig, V=V,\n",
    "            points=data_dict['points'][:, 1:].cpu().numpy(), #        \n",
    "            ref_boxes=pred_dicts[0]['pred_boxes'].cpu().numpy(),\n",
    "            ref_scores=pred_dicts[0]['pred_scores'].cpu().numpy(), ref_labels=pred_dicts[0]['pred_labels'].cpu().numpy()\n",
    "        )\n",
    "    # We visualize cloud overlayed by 3d boxes in a couple perspectives - *roughly* bev and front.\n",
    "    draw_bev = True\n",
    "    draw_front = True\n",
    "    if draw_bev:\n",
    "        mlab.view(azimuth=0, elevation=0, distance=30, focalpoint=(15,0,8))   #  ~ BEV\n",
    "        mlab.savefig('./rsc/orig_model_bev.png', figure=fig) \n",
    "    if draw_front:\n",
    "        mlab.view(azimuth=-179, elevation=70, distance=25, focalpoint=(5,0,5)) #  ~ front ..\n",
    "        mlab.savefig('./rsc/orig_model_front.png', figure=fig)\n",
    "else:\n",
    "    print(pred_dicts)  # not very readable, visualization is better:)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "471d9474",
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython\n",
    "width = 480  # use x4 (1920) for a full HD image\n",
    "IPython.display.Image('./rsc/orig_model_front.png', width=width)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67e9ddac",
   "metadata": {},
   "source": [
    "## Dissect the 2D parts and export to ONNX so we can parse for Hailo\n",
    "\n",
    "```\n",
    "We re-frame as pytorch module with tensor (not dictionary) I/O, starting from spatial_features and ending just before generate_boxes()  (so technically we don't map ALL of the dense-head, skipping the box generation; this can be remedied later with more semi-manual work representing via hailo-compatible ops, with anchors encoded as weights, etc etc.) \n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5758df5",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "\"\"\" Check out Bev_w_Head code in openpcdet2hailo_utils.py.\n",
    "    It's a slight rewrite of a section of the forward() method of openpcdet's pointpillars module, \n",
    "     desribing roughly the \"backbone_2d\" and \"dense_head\" parts. \n",
    "    Instead of using the common dictionary, it accepts the relevant tensor (\"spatial_features\") directly (vs. dictionary).\n",
    "    Wraps the original module which it accepts in constructor, code is copied from orig forward().\n",
    "\"\"\"\n",
    "reload(ohu)    \n",
    "bev_w_head = ohu.Bev_w_Head(model.backbone_2d, model.dense_head)\n",
    "export_res = torch.onnx.export(bev_w_head, args=(data_dict['spatial_features'],), \n",
    "                               f=\"pp_bev_w_head.onnx\", verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ae55450",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!onnxsim pointpillars_backbone2d.onnx pointpillars_backbone2d_simplified.onnx\n",
    "!onnxsim pp_bev_w_head.onnx pp_bev_w_head_simple.onnx"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98a5f37f",
   "metadata": {},
   "source": [
    "## Translate into Hailo internal representation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a12a2e20",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import onnxruntime\n",
    "from hailo_sdk_client import ClientRunner\n",
    "from hailo_sdk_common.targets.inference_targets import SdkNative\n",
    "from hailo_sdk_client import InferenceContext #SdkPartialNumeric, SdkNative # \n",
    "import tensorflow as tf\n",
    "import hailo_sdk_client, hailo_sdk_common\n",
    "print(hailo_sdk_client.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e78dea33",
   "metadata": {},
   "outputs": [],
   "source": [
    "runner = ClientRunner(hw_arch='hailo8')\n",
    "onnx_path = \"pp_bev_w_head_simple.onnx\"\n",
    "hn, npz = runner.translate_onnx_model(onnx_path)\n",
    "runner.save_har(har_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bfe3462e",
   "metadata": {},
   "source": [
    "## Verify inference equivalence with Hailo emulation injected into Pytorch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c6fd23a",
   "metadata": {},
   "source": [
    "```\n",
    "The HAR encapsulates the 2D part of the model, and we'd like to emulate it's execution by Hailo, as a part of the overall model. For that we'll be integrating via an alternate Pytorch model such that the 2D modules are replaced by a call into Hailo emulator executing the operations described by the HAR. To this end, we first create a torch submodule which executes hailo emulator, and can replace a section of the full torch model.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a29c9d8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Bev_W_Head_Hailo(torch.nn.Module):\n",
    "    \"\"\" Drop-in replacement to the sequence of original \"backbone-2d\" and \"dense_head\" modules, accepting and returning dictionary,\n",
    "        while under the hood using Hailo [emulator] implementation for the 2D CNN part, accepting and returning tensors I/O\n",
    "    \"\"\"\n",
    "    def __init__(self, runner, emulate_quantized=False, use_hw=False, generate_predicted_boxes=None):\n",
    "        super().__init__()\n",
    "        self._runner = runner\n",
    "        self.generate_predicted_boxes = generate_predicted_boxes\n",
    "        \n",
    "        if use_hw:\n",
    "            context_type = InferenceContext.SDK_HAILO_HW\n",
    "        elif emulate_quantized:\n",
    "            context_type = InferenceContext.SDK_QUANTIZED \n",
    "        else:\n",
    "            context_type = InferenceContext.SDK_FP_OPTIMIZED\n",
    "            \n",
    "        with runner.infer_context(context_type) as ctx:\n",
    "            self._hailo_model = runner.get_keras_model(ctx)   \n",
    "            \n",
    "    def forward(self, data_dict):        \n",
    "        spatial_features = data_dict['spatial_features']\n",
    "        \n",
    "        spatial_features_hailoinp = np.transpose(spatial_features.cpu().detach().numpy(), (0,2,3,1))\n",
    "        \n",
    "        # ============ Hailo-emulation of the Hailo-mapped part ==========\n",
    "        spatial_features_2d, cls_preds, box_preds, dir_cls_preds = \\\n",
    "                            self._hailo_model(spatial_features_hailoinp)\n",
    "        # ================================================================\n",
    "        \n",
    "        print('goo')\n",
    "        print(cls_preds.shape, type(cls_preds), box_preds.shape)\n",
    "        cls_preds = torch.Tensor(cls_preds.numpy()) # .permute(0, 2, 3, 1).contiguous()          # [N, H, W, C]\n",
    "        box_preds = torch.Tensor(box_preds.numpy()) # .permute(0, 2, 3, 1).contiguous()          # [N, H, W, C]\n",
    "        dir_cls_preds = torch.Tensor(dir_cls_preds.numpy()) # .permute(0, 2, 3, 1).contiguous()\n",
    "                \n",
    "        data_dict['spatial_features_2d'] = torch.Tensor(spatial_features_2d.numpy())\n",
    "        \n",
    "        batch_cls_preds, batch_box_preds = self.generate_predicted_boxes(\n",
    "            batch_size=data_dict['batch_size'],\n",
    "            cls_preds=cls_preds, box_preds=box_preds, dir_cls_preds=dir_cls_preds\n",
    "        )\n",
    "        data_dict['batch_cls_preds'] = batch_cls_preds\n",
    "        data_dict['batch_box_preds'] = batch_box_preds\n",
    "        data_dict['cls_preds_normalized'] = False\n",
    "\n",
    "        return data_dict\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8334873",
   "metadata": {},
   "source": [
    "```\n",
    "We'll test to verify predictions with offload to Hailo emulator similar to the original - \n",
    " that shows a successful parsing and full-precision emulation\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5f9623e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def quick_test(runner, hailoize=True, emulate_quantized=False, use_hw=False, fname='000102.bin', verbose=False):\n",
    "    \"\"\" Encapsulates a minimalistic test of the complete network with/without hailo offload emulation \n",
    "    \"\"\"\n",
    "    demo_dataset = ohu.DemoDataset(\n",
    "        dataset_cfg=cfg.DATA_CONFIG, class_names=cfg.CLASS_NAMES, training=False,\n",
    "        root_path=Path(fname), ext='.bin', logger=logger\n",
    "    )\n",
    "    logger.info(f'Total number of samples: \\t{len(demo_dataset)}')\n",
    "    \n",
    "    model_h = build_network(model_cfg=cfg.MODEL, num_class=len(cfg.CLASS_NAMES), dataset=demo_dataset)\n",
    "    model_h.load_params_from_file(filename=pth_name, logger=logger, to_cpu=True)\n",
    "    model_h.eval()\n",
    "\n",
    "    #bb2d_hailo1 = BB2d_Hailo(runner, pppost_onnx='./pp_tmp_post.onnx', emulate_quantized=emulate_quantized, use_hw=use_hw)\n",
    "    bev_w_head_hailo = Bev_W_Head_Hailo(runner, generate_predicted_boxes=model_h.dense_head.generate_predicted_boxes,\n",
    "                                        emulate_quantized=emulate_quantized, use_hw=use_hw)    \n",
    "    \n",
    "    if hailoize:        \n",
    "        # ==== Hook a call into Hailo by replacing parts of sequence by our rigged submodule ====\n",
    "        model_h.module_list = model_h.module_list[:2] + [bev_w_head_hailo]\n",
    "        # =======================================================================================                                          \n",
    "    \n",
    "    model_cpu = ohu.PointPillarsCPU(model_h) \n",
    "    logger.info(f'Total number of samples: \\t{len(demo_dataset)}')\n",
    "\n",
    "    for idx, data_dict in enumerate(demo_dataset):\n",
    "        logger.info(f'Visualized sample index: \\t{idx + 1}')\n",
    "        data_dict = demo_dataset.collate_batch([data_dict])            \n",
    "        pred_dicts, _ = model_cpu.forward(data_dict)\n",
    "        if verbose:\n",
    "            print(pred_dicts)\n",
    "        else:\n",
    "            print(pred_dicts[0]['pred_scores'][:7])\n",
    "            \n",
    "quick_test(runner, hailoize=False, emulate_quantized=False)     \n",
    "quick_test(runner, hailoize=True, emulate_quantized=False)     \n",
    "# This should give exact same result as we're yet to actually emulate the HW datapath,\n",
    "#   with its \"lossy-compression\" (e.g., 8b) features. \n",
    "# This will be possible after calibration and quantization of the model which will also enabling compilation for a physical HW."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35664f62",
   "metadata": {},
   "source": [
    "\n",
    "# Optimize & compile for Hailo HW\n",
    "\n",
    "## Make HW-compatible numeric model (q-HAR) \n",
    "```\n",
    " Model Optimization aka Quantization\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a46afbbd",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "### Create the input calibration dataset for the backbone_2d in numpy format \n",
    "```\n",
    "by collecting the spatial_features inputs to backbone_2d, over several point clouds\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ae2ae87",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO - provide download of at least 8 point clouds\n",
    "\n",
    "cfg_from_yaml_file_wrap(yaml_name, cfg)\n",
    "demo_dataset = ohu.DemoDataset(\n",
    "    dataset_cfg=cfg.DATA_CONFIG, class_names=cfg.CLASS_NAMES, training=False,\n",
    "    root_path=Path(sample_pointclouds), ext='.bin', logger=logger\n",
    ")\n",
    "logger.info(f'Total number of samples: \\t{len(demo_dataset)}')\n",
    "model = build_network(model_cfg=cfg.MODEL, num_class=len(cfg.CLASS_NAMES), dataset=demo_dataset)\n",
    "model.load_params_from_file(filename=pth_name, logger=logger, to_cpu=True)\n",
    "model.cuda()\n",
    "model.eval()\n",
    "cs_size = 8\n",
    "calib_set = np.zeros((cs_size, 496, 432, 64))\n",
    "np.set_printoptions(precision=2)\n",
    "\n",
    "perc4stats = [50, 90, 98.6, 99.7, 99.9]\n",
    "with torch.no_grad():\n",
    "    for idx, _data_dict in enumerate(demo_dataset):        \n",
    "        print(f\"cloud #{idx}\")\n",
    "        if idx >= cs_size:\n",
    "            break\n",
    "        _data_dict = demo_dataset.collate_batch([_data_dict])\n",
    "        load_data_to_gpu(_data_dict)\n",
    "        pred_dicts, goo = model.forward(_data_dict)       \n",
    "        calib_set[idx] = np.transpose(_data_dict['spatial_features'].cpu().numpy(), (0,2,3,1))\n",
    "        # Basic stats just to verify there's some data diversity (just in top percentile apparently...)\n",
    "        print(f'basic stats - percentile {perc4stats} of data (@ 2d-net input)', \\\n",
    "              np.percentile((np.abs(calib_set[idx])), perc4stats))\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "374681ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data is sparse, roughly ~ 99% zeros. But our quantization doesn't seem to be affected.\n",
    "np.save('calib.npy', calib_set)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2060657",
   "metadata": {},
   "source": [
    "### Run the model optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc3869ea",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "do_optimize = True\n",
    "q_har_name = har_name.replace('.har', '.q.har')\n",
    "\n",
    "if do_optimize:\n",
    "    runner.optimize(calib_set)\n",
    "    runner.save_har(q_har_name)\n",
    "else:\n",
    "    runner = ClientRunner(har_path=q_har_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "392cbb76",
   "metadata": {},
   "source": [
    "### More emulation testing - now with the quantized model, closely following the numerics of the hardware\n",
    "```\n",
    " Hopefully, we see here results with small-ish differences\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "815e9188",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "quick_test(runner, hailoize=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f3bd9c8",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "quick_test(runner, hailoize=True, emulate_quantized=True)\n",
    "\n",
    "# Hopefully, we see here results with small-ish differences"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2dc4ba82",
   "metadata": {},
   "source": [
    "## Make optimal physical allocation and build the binary (HEF)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d84354f7",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "do_compile = True\n",
    "if do_compile:\n",
    "    alls_line1 = 'shortcut_concat1_conv20 = shortcut(concat1, conv20)\\n'\n",
    "    open('helper.alls','w').write(alls_line1)  #   !!!!\n",
    "    runner.load_model_script('./helper.alls') \n",
    "\n",
    "    compiled_model=runner.compile()    \n",
    "    open(hef_name, 'wb').write(compiled_model)\n",
    "\n",
    "# Expected results:\n",
    "# [info] | Cluster   | Control Utilization | Compute Utilization | Memory Utilization |\n",
    "# [info] +-----------+---------------------+---------------------+--------------------+\n",
    "#                                        ...\n",
    "# [info] +-----------+---------------------+---------------------+--------------------+\n",
    "# [info] | Total     | 62.5%               | 63.7%               | 37.2%              |"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cd60dfd",
   "metadata": {},
   "source": [
    "# Running end2end with 2D part offloaded to Hailo HW"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2309458",
   "metadata": {},
   "source": [
    "```\n",
    "We will be using HailoRT asynchronous send/receive to demonstrate readiness for a fully pipelined processing.\n",
    "\n",
    "Therefore, we integrate the Torch and Hailo parts differently from what we did for the emulation testing. We rip off some code from original forward() to build two torch models which encapsulate everything that happens before and after the module offloaded to Hailo.  See the code for 'PP_Pre_Bev_w_Head', 'PP_Post_Bev_w_Head' classes in the accompanying utils file.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b2360de",
   "metadata": {},
   "outputs": [],
   "source": [
    "reload(ohu);  # PP_Pre_Bev_w_Head, PP_Post_Bev_w_Head, "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "260436cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" ================================================\n",
    "    Standard wrapping of HailoRT API with send/receive processes, \n",
    "    similar to other Hailo examples (or HailoRT tutorial)\n",
    "    ================================================\n",
    "\"\"\"\n",
    "from multiprocessing import Process, Queue\n",
    "from hailo_platform import (HEF, PcieDevice, VDevice, HailoStreamInterface, ConfigureParams,\n",
    " InputVStreamParams, OutputVStreamParams, InputVStreams, OutputVStreams, FormatType)\n",
    "\n",
    "def send_from_queue(configured_network, read_q, num_images, start_time):\n",
    "    \"\"\" Bridging a queue into Hailo platform FEED. To be run as a separate process. \n",
    "        Reads (preprocessed) images from a given queue, and sends them serially to Hailo platform.        \n",
    "    \"\"\"    \n",
    "    configured_network.wait_for_activation(1000)\n",
    "    vstreams_params = InputVStreamParams.make(configured_network, quantized=False, format_type=FormatType.FLOAT32)\n",
    "    print('Starting sending input images to HW inference...\\n')\n",
    "    with InputVStreams(configured_network, vstreams_params) as vstreams:\n",
    "        vstream_to_buffer = {vstream: np.ndarray([1] + list(vstream.shape), dtype=vstream.dtype) for vstream in vstreams}\n",
    "        for i in range(num_images):\n",
    "            hailo_inp = read_q.get()\n",
    "            for vstream, _ in vstream_to_buffer.items():                                \n",
    "                vstream.send(hailo_inp)\n",
    "            print(f'sent img #{i}')\n",
    "    print(F'Finished send after {(time.time()-start_time) :.1f}')\n",
    "    return 0\n",
    "\n",
    "def recv_to_queue(configured_network, write_q, num_images, start_time):\n",
    "    \"\"\" Bridging Hailo platform OUTPUT into a queue. To be run as a separate process. \n",
    "        Reads output data from Hailo platform and sends them serially to a given queue.\n",
    "    \"\"\"\n",
    "    configured_network.wait_for_activation(1000)\n",
    "    vstreams_params = OutputVStreamParams.make_from_network_group(configured_network, quantized=False, format_type=FormatType.FLOAT32)\n",
    "    print('Starting receving HW inference output..\\n')\n",
    "    with OutputVStreams(configured_network, vstreams_params) as vstreams:\n",
    "        # print('vstreams_params', vstreams_params)\n",
    "        for i in range(num_images):            \n",
    "            hailo_out = {vstream.name: np.expand_dims(vstream.recv(), 0) for vstream in vstreams}                          \n",
    "            write_q.put(hailo_out)\n",
    "            print(f'received img #{i}')\n",
    "    print(F'Finished recv after {time.time()-start_time :.1f}')\n",
    "    return 0\n",
    "\n",
    "\"\"\" ==================================\n",
    "    Some final wrapping of pre and post \n",
    "    ==================================\n",
    "\"\"\"\n",
    "def generate_data_dicts(demo_dataset, num_images, pp_pre_bev_w_head):    \n",
    "    for idx, data_dict in enumerate(demo_dataset):\n",
    "        if idx > num_images :\n",
    "            break                        \n",
    "        data_dict = demo_dataset.collate_batch([data_dict])        \n",
    "        ohu.load_data_to_CPU(data_dict)\n",
    "        # ------ (!) Applying torch PRE-processing -------\n",
    "        data_dict = pp_pre_bev_w_head.forward(data_dict)        \n",
    "        # ------------------------------------------------\n",
    "        logger.info(f'preprocessed sample #{idx}')\n",
    "        yield data_dict\n",
    "        \n",
    "def generate_hailo_inputs(demo_dataset, num_images, pp_pre_bev_w_head):\n",
    "    \"\"\" generator-style encapsulation for preprocessing inputs for Hailo HW feed\n",
    "    \"\"\"\n",
    "    for data_dict in generate_data_dicts(demo_dataset, num_images, pp_pre_bev_w_head):\n",
    "        spatial_features = data_dict['spatial_features']\n",
    "        spatial_features_hailoinp = np.transpose(spatial_features.cpu().detach().numpy(), (0,2,3,1))        \n",
    "        yield data_dict, spatial_features_hailoinp\n",
    "\n",
    "def post_proc_from_queue(recv_queue, num_images, pp_post_bev_w_head,\n",
    "                         output_layers_order = ['model/concat1', 'model/conv18', 'model/conv19', 'model/conv20']):\n",
    "    results = []\n",
    "    for i in range(num_images):\n",
    "        print('a')\n",
    "        t_ = time.time()                \n",
    "        while(recv_queue.empty() and time.time()-t_ <3):                    \n",
    "            time.sleep(0.01)\n",
    "        if recv_queue.empty():\n",
    "            print(\"RECEIVE TIMEOUT!\")\n",
    "            break\n",
    "        hailo_out = recv_queue.get(0)   \n",
    "        if False:\n",
    "            for k, ho in hailo_out.items():\n",
    "                print(k, \"ho.shape\", ho.shape)\n",
    "\n",
    "        bev_out = (hailo_out[lname] for lname in output_layers_order)        \n",
    "        \n",
    "        # ------ (!) Applying torch POST-processing -------\n",
    "        pred_dicts, _ = pp_post_bev_w_head(bev_out)\n",
    "        # ------------------------------------------------\n",
    "        results.append(pred_dicts)\n",
    "        # print(pred_dicts[0]['pred_scores'])  \n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "daf99685",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import time, onnxruntime\n",
    "\n",
    "data_source = demo_pointcloud  # replace by a folder for a more serious test\n",
    "num_images = 1\n",
    "\n",
    "cfg_from_yaml_file_wrap(yaml_name, cfg)\n",
    "logger = common_utils.create_logger()\n",
    "demo_dataset = ohu.DemoDataset(\n",
    "    dataset_cfg=cfg.DATA_CONFIG, class_names=cfg.CLASS_NAMES, training=False,\n",
    "    root_path=Path(data_source), ext='.bin', logger=logger\n",
    "    )\n",
    "model = get_model(cfg, pth_name, demo_dataset)\n",
    "\n",
    "# Library creates the anchors in cuda by defalt (applying .cuda() in internal implementation)\n",
    "model.dense_head.anchors = [anc.cpu() for anc in model.dense_head.anchors]\n",
    "\n",
    "\"\"\" (!) Slicing off the torch model all that happens before and after Hailo\n",
    "\"\"\"\n",
    "pp_pre_bev_w_head = ohu.PP_Pre_Bev_w_Head(model)\n",
    "pp_post_bev_w_head = ohu.PP_Post_Bev_w_Head(model)\n",
    "    \n",
    "with VDevice() as target:\n",
    "    hef = HEF(hef_name)\n",
    "    configure_params = ConfigureParams.create_from_hef(hef, interface=HailoStreamInterface.PCIe)\n",
    "    network_group = target.configure(hef, configure_params)[0]\n",
    "    network_group_params = network_group.create_params()\n",
    "    recv_queue = Queue()\n",
    "    send_queue = Queue()            \n",
    "    start_time = time.time()\n",
    "    results = []\n",
    "    hw_send_process = Process(target=send_from_queue, args=(network_group, send_queue, num_images, start_time))\n",
    "    hw_recv_process = Process(target=recv_to_queue, args=(network_group, recv_queue, num_images, start_time))            \n",
    "    \n",
    "    with network_group.activate(network_group_params): \n",
    "        hw_recv_process.start()        \n",
    "        hw_send_process.start()\n",
    "\n",
    "        # Note: still short of 'infinitely streamed' demo as neither (A) nor (B) can offload to a separate process \n",
    "        #       because of unfortunate cuda issues with multiproc (cuda not used in (A) but still loaded).\n",
    "        # One partial solution could be to intermittently feed-in/process-out according to queues status, keeping HW busy\n",
    "        \n",
    "        # (A)\n",
    "        for data_dict, hailo_inp in generate_hailo_inputs(demo_dataset, num_images, pp_pre_bev_w_head):\n",
    "            send_queue.put(hailo_inp)                    \n",
    "\n",
    "        # (B)\n",
    "        results = post_proc_from_queue(recv_queue, num_images, pp_post_bev_w_head)\n",
    "                             \n",
    "hw_recv_process.join(10)\n",
    "hw_send_process.join(10)\n",
    "\n",
    "pred_dicts = results[-1]\n",
    "print(pred_dicts[0]['pred_scores'])   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afd5f3ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "if do_visualize: \n",
    "    fig = mlab.figure(size=(1920, 1080))\n",
    "\n",
    "    s = ohu.my_draw_scenes(fig=fig, V=V,\n",
    "            points=data_dict['points'][:, 1:].detach(),       \n",
    "            ref_boxes=pred_dicts[0]['pred_boxes'].detach(),\n",
    "            ref_scores=pred_dicts[0]['pred_scores'].detach(), ref_labels=pred_dicts[0]['pred_labels'].detach()\n",
    "        )\n",
    "    \n",
    "    draw_bev = True\n",
    "    draw_front = True\n",
    "    if draw_bev:\n",
    "        mlab.view(azimuth=0, elevation=0, distance=30, focalpoint=(15,0,8))   #  ~ BEV\n",
    "        mlab.savefig('./rsc/hailo_assisted_model_bev.png', figure=fig) \n",
    "    if draw_front:\n",
    "        mlab.view(azimuth=-179, elevation=70, distance=25, focalpoint=(5,0,5)) #  ~ front ..\n",
    "        mlab.savefig('./rsc/hailo_assisted_model_front.png', figure=fig)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08d0f8e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython\n",
    "width = 480  # use x4 (1920) for a full HD image\n",
    "IPython.display.Image('./rsc/hailo_assisted_model_front.png', width=width)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03d83d47",
   "metadata": {},
   "source": [
    "# Appendix A:   OpenPCDet integration details\n",
    "\n",
    "```\n",
    "Here we explain how we arrived at the openpcdet-hailo integration details that might have appeared as \"magic\" in the code above and in the integration-utils module - which is by and large simply pieces of opencpcdet code copied and modified in specific ways to make integration convenient. For example, how to isolate the torch subgraph to translate to Hailo, patch the inference func with a call into Hailo emulator/hardware, etc.\n",
    "```\n",
    "\n",
    "**This might be useful in case you want to integrate and Hailo-offload a NEW OpenPCDet model**\n",
    "\n",
    "**In that case, you might want to repeat the exploration steps below and modify the integration code accordingly.**\n",
    "\n",
    "```\n",
    "NOTE: apart from inspecting tensor structure on model_cpu object (that we assume to have been run during \"sanity test\" above), we will also be just statically eyeballing some openpcdet code via python's \"inspecting\". This might look cumbersome, and sure you can just browse the repo. However, from our experience, the \"inspect\" approach is not only convenient for display in notebook, but actually useful to quickly get to the code that **actually gets to run**, bypassing all configuration switches which might be rather complex in big modular 'many-model' repos.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df6ad290",
   "metadata": {},
   "source": [
    "## Explore model structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff97f042",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import inspect\n",
    "print(inspect.getsource(model.forward))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae7633a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.module_list[0] == model.vfe, model.module_list[1] == model.map_to_bev_module, \\\n",
    "  model.module_list[2] == model.backbone_2d, model.module_list[3] == model.dense_head"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0e3bb5b",
   "metadata": {},
   "source": [
    "```\n",
    "From this [..verification of an educated guess..] we learn that model is a sequence of 4 modules.\n",
    "We want to offload the following two:\n",
    "A. model.module_list[2] == model.backbone_2d , along with\n",
    "B. model.module_list[3] == model.dense_head.\n",
    "\n",
    "Now let's take apart the modules and their forward() and I/O:\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98b03e63",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for submodule in model.module_list:\n",
    "    for ii in range(3):\n",
    "        print('='*100)\n",
    "    print(submodule)\n",
    "    print(inspect.getsource(submodule.forward))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d28239a9",
   "metadata": {},
   "source": [
    "``` \n",
    "Also let's take a look at shapes of all intermediate results\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9452765",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Assuming live model_cpu and data_dict from above.. (e.g. sanity test has been run)\n",
    "\n",
    "model_cpu.eval()\n",
    "pred_dicts, _ = model_cpu(data_dict)\n",
    "# .. we see that in openpcdet the forward() graudally populates data_dict with all intermediate results\n",
    "\n",
    "for k,v in data_dict.items():\n",
    "    try: \n",
    "        print(k, v.shape) #, np.max(np.abs(v.cpu().numpy())))\n",
    "    except:\n",
    "        print(k, v)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "488ec441",
   "metadata": {},
   "source": [
    "```\n",
    "From this we learn that input to 2d-backbone part is the \"spatial_features\" key,\n",
    "   tensor shaped as torch.Size([1, 64, 496, 432]\n",
    "\n",
    "And the output from the \"dense\" head module is 'batch_cls_preds' and 'batch_box_preds' keys.\n",
    "\n",
    "Let's explore the last non-neural part of the dense_head - box decoder from anchors and net results, encapsulated in generate_predicted_boxes(). This part will take some further tweaks to get mapped to Hailo, as exists for 2D detection heads (ssd, yolo, etc.).  We leave that for future work\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85944ed4",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(model.dense_head.conv_dir_cls)\n",
    "print(inspect.getsource(model.dense_head.generate_predicted_boxes))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef599699",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "\"\"\" Yet more torch inspection APIs ... \"\"\" \n",
    "\n",
    "print_backbone_2d_details = False\n",
    "print_all_modules_details = False\n",
    "print_named_children = False\n",
    "\n",
    "if print_backbone_2d_details:\n",
    "    print(model.backbone_2d)\n",
    "if print_all_modules_details:\n",
    "    for ml in model.module_list:\n",
    "        print('='*50)\n",
    "        print(ml)\n",
    "if print_named_children:\n",
    "    for c in model.named_children():\n",
    "        print(c)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
